package com.young.mbgdiy.plugin;


import org.mybatis.generator.api.IntrospectedColumn;
import org.mybatis.generator.api.IntrospectedTable;
import org.mybatis.generator.api.PluginAdapter;
import org.mybatis.generator.api.dom.xml.Attribute;
import org.mybatis.generator.api.dom.xml.XmlElement;

import java.util.List;

/**
 * XML insert 标签添加自定义属性
 *
 * @author ：<a href="mailto:youngkun2016@163.com">young</a>
 * @date ：Created in 2020/4/5
 */
public class PrimaryKeyMapperPlugin extends PluginAdapter {

    private static final String KEYCOLUMN = "keyColumn";
    private static final String KEYPROPERTY = "keyProperty";

    @Override
    public boolean validate(List<String> warnings) {
        return true;
    }

    @Override
    public boolean sqlMapInsertSelectiveElementGenerated(XmlElement element,
                                                         IntrospectedTable introspectedTable) {
        return addKeyProperty(element, introspectedTable);
    }

    @Override
    public boolean sqlMapInsertElementGenerated(XmlElement element,
                                                IntrospectedTable introspectedTable) {
        return addKeyProperty(element, introspectedTable);
    }

    /**
     * 添加属性
     * keyColumn
     * keyProperty
     *
     * @param element
     * @param introspectedTable
     * @return
     */
    private boolean addKeyProperty(XmlElement element,
                                   IntrospectedTable introspectedTable) {
        List<IntrospectedColumn> keyColumns = introspectedTable.getPrimaryKeyColumns();
        // 针对单一主键情况
        if (keyColumns.size() == 1) {
            // 获取主键列
            IntrospectedColumn keyColumn = keyColumns.get(0);
            // 数据库表中主键field名称
            element.addAttribute(new Attribute(KEYCOLUMN, keyColumn.getActualColumnName()));
            // 主键对应实体类中的filed
            element.addAttribute(new Attribute(KEYPROPERTY, keyColumn.getJavaProperty()));
        }
        return true;
    }
}
